import cv2
import numpy as np
import os

class MangaPanelExtractor:
    def __init__(self):
        self.white_threshold = 220  # 白色区域的阈值
        self.min_line_width = 2     # 最小线宽
        self.min_panel_size = 50    # 最小分镜大小
        self.debug = True

    def find_white_lines(self, image):
        """查找白色分割线"""
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用高斯模糊减少噪点
        blurred = cv2.GaussianBlur(gray, (3, 3), 0)
        
        # 二值化
        _, binary = cv2.threshold(blurred, self.white_threshold, 255, cv2.THRESH_BINARY)
        
        # 获取图像尺寸
        height, width = binary.shape
        
        # 查找水平和垂直投影
        h_projection = np.sum(binary, axis=1) / 255  # 水平投影
        v_projection = np.sum(binary, axis=0) / 255  # 垂直投影
        
        # 查找水平分割线
        h_lines = []
        start_y = None
        for y in range(height):
            if h_projection[y] > width * 0.5:  # 如果超过50%是白色
                if start_y is None:
                    start_y = y
            elif start_y is not None:
                if y - start_y >= self.min_line_width:
                    h_lines.append((start_y, y))
                start_y = None
        
        # 查找垂直分割线
        v_lines = []
        start_x = None
        for x in range(width):
            if v_projection[x] > height * 0.5:  # 如果超过50%是白色
                if start_x is None:
                    start_x = x
            elif start_x is not None:
                if x - start_x >= self.min_line_width:
                    v_lines.append((start_x, x))
                start_x = None
        
        return h_lines, v_lines

    def extract_panels(self, image_path, output_dir):
        """提取漫画分镜"""
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        print(f"\n开始处理图片: {image_path}")
        
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图片: {image_path}")
        
        # 获取图片尺寸
        height, width = image.shape[:2]
        print(f"图片尺寸: {width}x{height}")
        
        # 查找白色分割线
        h_lines, v_lines = self.find_white_lines(image)
        
        # 在调试图像上绘制分割线
        debug_image = image.copy()
        
        # 绘制水平分割线
        for y1, y2 in h_lines:
            cv2.line(debug_image, (0, y1), (width, y1), (0, 255, 0), 2)
            cv2.line(debug_image, (0, y2), (width, y2), (0, 255, 0), 2)
        
        # 绘制垂直分割线
        for x1, x2 in v_lines:
            cv2.line(debug_image, (x1, 0), (x1, height), (0, 255, 0), 2)
            cv2.line(debug_image, (x2, 0), (x2, height), (0, 255, 0), 2)
        
        # 保存调试图像
        cv2.imwrite(os.path.join(output_dir, "detected_lines.jpg"), debug_image)
        
        # 确保有分割线
        if not h_lines:
            h_lines = [(0, height)]
        if not v_lines:
            v_lines = [(0, width)]
        
        # 添加图片边界
        if h_lines[0][0] > 0:
            h_lines.insert(0, (0, 0))
        if h_lines[-1][1] < height:
            h_lines.append((height, height))
        
        if v_lines[0][0] > 0:
            v_lines.insert(0, (0, 0))
        if v_lines[-1][1] < width:
            v_lines.append((width, width))
        
        # 提取分镜
        panel_count = 0
        for i in range(len(h_lines)-1):
            for j in range(len(v_lines)-1):
                y1 = h_lines[i][1]
                y2 = h_lines[i+1][0]
                x1 = v_lines[j][1]
                x2 = v_lines[j+1][0]
                
                # 检查区域大小
                if (x2 - x1) < self.min_panel_size or (y2 - y1) < self.min_panel_size:
                    continue
                
                # 提取分镜
                panel = image[y1:y2, x1:x2]
                
                # 检查是否为空白区域
                gray_panel = cv2.cvtColor(panel, cv2.COLOR_BGR2GRAY)
                if np.mean(gray_panel) > self.white_threshold:
                    continue
                
                panel_count += 1
                output_path = os.path.join(output_dir, f"panel_{panel_count}.jpg")
                cv2.imwrite(output_path, panel)
                print(f"保存分镜 {panel_count}: {output_path}")
        
        return panel_count

def main():
    # 创建实例
    extractor = MangaPanelExtractor()
    
    # 确保输入和输出目录存在
    input_dir = "pic"
    output_base_dir = "extracted_panels"
    
    if not os.path.exists(input_dir):
        os.makedirs(input_dir)
        print(f"创建输入目录: {input_dir}")
    
    if not os.path.exists(output_base_dir):
        os.makedirs(output_base_dir)
        print(f"创建输出目录: {output_base_dir}")
    
    # 获取所有JPG文件
    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.jpg')]
    
    print(f"\n找到 {len(image_files)} 个图片文件")
    
    if not image_files:
        print(f"\n错误：在 {input_dir} 目录中没有找到图片文件！")
        return
    
    # 处理每个图片文件
    for image_file in image_files:
        print(f"\n处理文件: {image_file}")
        input_path = os.path.join(input_dir, image_file)
        
        # 创建输出目录
        output_dir = os.path.join(output_base_dir, os.path.splitext(image_file)[0])
        
        try:
            # 如果输出目录已存在，先删除
            if os.path.exists(output_dir):
                import shutil
                shutil.rmtree(output_dir)
                print(f"清理已存在的输出目录: {output_dir}")
            
            # 处理图片
            num_panels = extractor.extract_panels(input_path, output_dir)
            print(f"成功从 {image_file} 中提取了 {num_panels} 个分镜到 {output_dir}")
            
        except Exception as e:
            print(f"处理 {image_file} 时出错: {str(e)}")
            continue

if __name__ == "__main__":
    main() 